import cv2
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.transforms import Compose

from .dinov2 import DINOv2
from .util.blocks import FeatureFusionBlock, _make_scratch, ControlFeatureFusionBlock
from .util.transform import sample2tensor
from .dpt import _make_fusion_block
from PIL import Image

    
def _make_control_fusion_block(features, use_bn, init_control_zero=True, size=None):
    return ControlFeatureFusionBlock(
        features,
        nn.ReLU(False),
        deconv=False,
        bn=use_bn,
        expand=False,
        align_corners=True,
        size=size,
        init_control_zero=init_control_zero,
        
    )

class MainDPTHead(nn.Module):
    def __init__(
        self, 
        in_channels, 
        features=256, 
        use_bn=False, 
        out_channels=[256, 512, 1024, 1024], 
        use_clstoken=False
    ):
        super(MainDPTHead, self).__init__()        
        self.use_clstoken = use_clstoken
        
        self.projects = nn.ModuleList([
            nn.Conv2d(
                in_channels=in_channels,
                out_channels=out_channel,
                kernel_size=1,
                stride=1,
                padding=0,
            ) for out_channel in out_channels
        ])
        
        self.resize_layers = nn.ModuleList([
            nn.ConvTranspose2d(
                in_channels=out_channels[0],
                out_channels=out_channels[0],
                kernel_size=4,
                stride=4,
                padding=0),
            nn.ConvTranspose2d(
                in_channels=out_channels[1],
                out_channels=out_channels[1],
                kernel_size=2,
                stride=2,
                padding=0),
            nn.Identity(),
            nn.Conv2d(
                in_channels=out_channels[3],
                out_channels=out_channels[3],
                kernel_size=3,
                stride=2,
                padding=1)
        ])
        
        if use_clstoken:
            self.readout_projects = nn.ModuleList()
            for _ in range(len(self.projects)):
                self.readout_projects.append(
                    nn.Sequential(
                        nn.Linear(2 * in_channels, in_channels),
                        nn.GELU()))
        
        self.scratch = _make_scratch(
            out_channels,
            features,
            groups=1,
            expand=False,
        )
        
        self.scratch.stem_transpose = None
        
        self.scratch.refinenet1 = _make_control_fusion_block(features, use_bn, True)
        self.scratch.refinenet2 = _make_control_fusion_block(features, use_bn, True)
        self.scratch.refinenet3 = _make_control_fusion_block(features, use_bn, True)
        self.scratch.refinenet4 = _make_control_fusion_block(features, use_bn, True)
        
        head_features_1 = features
        head_features_2 = 32
        
        self.scratch.output_conv1 = nn.Conv2d(head_features_1, head_features_1 // 2, kernel_size=3, stride=1, padding=1)
        self.scratch.output_conv2 = nn.Sequential(
            nn.Conv2d(head_features_1 // 2, head_features_2, kernel_size=3, stride=1, padding=1),
            nn.ReLU(True),
            nn.Conv2d(head_features_2, 1, kernel_size=1, stride=1, padding=0),
            nn.ReLU(True),
            nn.Identity(),
        )
    
    def forward(self, out_features, control_layers, patch_h, patch_w):
        out = []
        for i, x in enumerate(out_features):
            if self.use_clstoken:
                x, cls_token = x[0], x[1]
                readout = cls_token.unsqueeze(1).expand_as(x)
                x = self.readout_projects[i](torch.cat((x, readout), -1))
            else:
                x = x[0]
            
            x = x.permute(0, 2, 1).reshape((x.shape[0], x.shape[-1], patch_h, patch_w))
            
            x = self.projects[i](x)
            x = self.resize_layers[i](x)
            
            out.append(x)
        
        layer_1, layer_2, layer_3, layer_4 = out
        
        layer_1_rn = self.scratch.layer1_rn(layer_1)
        layer_2_rn = self.scratch.layer2_rn(layer_2)
        layer_3_rn = self.scratch.layer3_rn(layer_3)
        layer_4_rn = self.scratch.layer4_rn(layer_4)
        
        path_4 = self.scratch.refinenet4(layer_4_rn, None, control_layers[3], size=layer_3_rn.shape[2:])        
        path_3 = self.scratch.refinenet3(path_4, layer_3_rn, control_layers[2], size=layer_2_rn.shape[2:])
        path_2 = self.scratch.refinenet2(path_3, layer_2_rn, control_layers[1], size=layer_1_rn.shape[2:])
        path_1 = self.scratch.refinenet1(path_2, layer_1_rn, control_layers[0])
        
        out = self.scratch.output_conv1(path_1)
        out = F.interpolate(out, (int(patch_h * 14), int(patch_w * 14)), mode="bilinear", align_corners=True)
        out = self.scratch.output_conv2(out)
        
        return out
    
    def set_control_grad(self):
        self.scratch.refinenet1.set_control_grad()
        self.scratch.refinenet2.set_control_grad()
        self.scratch.refinenet3.set_control_grad()
        self.scratch.refinenet4.set_control_grad()
        
    def get_control_params(self):
        return self.scratch.refinenet1.get_control_params() + self.scratch.refinenet2.get_control_params() + self.scratch.refinenet3.get_control_params() + self.scratch.refinenet4.get_control_params()
    
class ControlDPTHead(nn.Module):
    def __init__(
        self, 
        in_channels, 
        features=256, 
        use_bn=False, 
        out_channels=[256, 512, 1024, 1024], 
        use_clstoken=False
    ):
        super(ControlDPTHead, self).__init__()        
        self.use_clstoken = use_clstoken
        
        self.projects = nn.ModuleList([
            nn.Conv2d(
                in_channels=in_channels,
                out_channels=out_channel,
                kernel_size=1,
                stride=1,
                padding=0,
            ) for out_channel in out_channels
        ])
        
        self.resize_layers = nn.ModuleList([
            nn.ConvTranspose2d(
                in_channels=out_channels[0],
                out_channels=out_channels[0],
                kernel_size=4,
                stride=4,
                padding=0),
            nn.ConvTranspose2d(
                in_channels=out_channels[1],
                out_channels=out_channels[1],
                kernel_size=2,
                stride=2,
                padding=0),
            nn.Identity(),
            nn.Conv2d(
                in_channels=out_channels[3],
                out_channels=out_channels[3],
                kernel_size=3,
                stride=2,
                padding=1)
        ])
        
        if use_clstoken:
            self.readout_projects = nn.ModuleList()
            for _ in range(len(self.projects)):
                self.readout_projects.append(
                    nn.Sequential(
                        nn.Linear(2 * in_channels, in_channels),
                        nn.GELU()))
        
        self.scratch = _make_scratch(
            out_channels,
            features,
            groups=1,
            expand=False,
        )
        
        self.scratch.stem_transpose = None
    
    def forward(self, out_features, patch_h, patch_w):
        out = []
        for i, x in enumerate(out_features):
            if self.use_clstoken:
                x, cls_token = x[0], x[1]
                readout = cls_token.unsqueeze(1).expand_as(x)
                x = self.readout_projects[i](torch.cat((x, readout), -1))
            else:
                x = x[0]
            
            x = x.permute(0, 2, 1).reshape((x.shape[0], x.shape[-1], patch_h, patch_w))
            
            x = self.projects[i](x)
            x = self.resize_layers[i](x)
            
            out.append(x)
        
        layer_1, layer_2, layer_3, layer_4 = out
        
        layer_1_rn = self.scratch.layer1_rn(layer_1)
        layer_2_rn = self.scratch.layer2_rn(layer_2)
        layer_3_rn = self.scratch.layer3_rn(layer_3)
        layer_4_rn = self.scratch.layer4_rn(layer_4)

        out = [layer_1_rn, layer_2_rn, layer_3_rn, layer_4_rn]
        
        return out

# 添加了Control的DepthAnythingV2
class ControlDepthAnythingV2(nn.Module):
    def __init__(
        self, 
        encoder='vitl', 
        features=256, 
        out_channels=[256, 512, 1024, 1024], 
        use_bn=False, 
        use_clstoken=False,
        max_depth=20.0
    ):
        super(ControlDepthAnythingV2, self).__init__()
        
        self.intermediate_layer_idx = {
            'vits': [2, 5, 8, 11],
            'vitb': [2, 5, 8, 11], 
            'vitl': [4, 11, 17, 23], 
            'vitg': [9, 19, 29, 39]
        }
        
        self.max_depth = max_depth
        
        self.encoder = encoder
        self.pretrained = DINOv2(model_name=encoder)
        self.control_encoder = DINOv2(model_name=encoder)
        
        self.depth_head = MainDPTHead(self.pretrained.embed_dim, features, use_bn, out_channels=out_channels, use_clstoken=use_clstoken)
        self.control_depth_head = ControlDPTHead(self.control_encoder.embed_dim, features, use_bn, out_channels=out_channels, use_clstoken=use_clstoken)
    
    # x: image, condition: control (depth, depth, mask)
    def forward(self, x, condition):
        patch_h, patch_w = x.shape[-2] // 14, x.shape[-1] // 14
        features = self.pretrained.get_intermediate_layers(x, self.intermediate_layer_idx[self.encoder], return_class_token=True)
        control_features = self.control_encoder.get_intermediate_layers(condition, self.intermediate_layer_idx[self.encoder], return_class_token=True)
        
        control_layers = self.control_depth_head(control_features, patch_h, patch_w)
        depth = self.depth_head(features, control_layers, patch_h, patch_w)
        depth = F.relu(depth)
        
        return depth.squeeze(1)
    
    # raw_image: RBG(ndarray), raw_disparity: disparity(ndarray), mask: mask(ndarray)
    @torch.no_grad()
    def infer_image(self, raw_image, raw_disparity, mask, input_size=518):
        if isinstance(raw_image, np.ndarray):
            raw_image = torch.from_numpy(raw_image)
        if isinstance(raw_disparity, np.ndarray):
            raw_disparity = torch.from_numpy(raw_disparity)
        if isinstance(mask, np.ndarray):
            mask = torch.from_numpy(mask)
        # 添加batch维度
        if len(raw_image.shape) == 3:
            raw_image = raw_image.unsqueeze(0)
        if len(raw_disparity.shape) == 2:
            raw_disparity = raw_disparity.unsqueeze(0)
        if len(mask.shape) == 3:
            mask = mask.unsqueeze(0)
        
        raw_disparity = raw_disparity.unsqueeze(-1).repeat(1, 1, 1, 3)
        
        h, w = raw_image.shape[1], raw_image.shape[2]
        samples = sample2tensor(raw_image, raw_disparity, mask, input_size)
        
        DEVICE = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
        raw_disparity = samples['disparity'].to(DEVICE)
        image = samples['image'].to(DEVICE)
        mask = samples['mask'].to(DEVICE)
        # [depth[c0], depth[c1], mask[c2]]
        condition = torch.concat([raw_disparity[...,:2], mask[...,2:]], dim=-1)
        condition = condition.permute(0, 3, 1, 2)
        image = image.permute(0, 3, 1, 2)
        depth = self.forward(image, condition)
        
        depth = F.interpolate(depth[:, None], (h, w), mode="bilinear", align_corners=True)[0, 0]
        
        return depth.cpu().numpy()
    
    
    # 设置control的参数grad为true，其他为false
    def set_grad(self):
        
        for param in self.depth_head.parameters():
            param.requires_grad = False
        for param in self.pretrained.parameters():
            param.requires_grad = False
        for param in self.control_encoder.parameters():
            param.requires_grad = True
        for param in self.control_depth_head.parameters():
            param.requires_grad = True
            
        self.depth_head.set_control_grad()
    
    # 获取和control有关的参数
    def get_control_params(self):
        return list(self.control_encoder.parameters()) + list(self.control_depth_head.parameters()) + self.depth_head.get_control_params()
        
    def init_by_main_state_dict(self, state_dict):
        missing_keys, unexpected_keys = self.load_state_dict(state_dict, strict=False)
        # if missing_keys:
        #     print(f"Main Missing keys: {missing_keys}")
        # if unexpected_keys:
        #     print(f"Main Unexpected keys: {unexpected_keys}")
        missing_keys, unexpected_keys = None, None
        missing_keys, unexpected_keys = self.control_encoder.load_state_dict(self.pretrained.state_dict(), strict=False)
        if missing_keys:
            print(f"Control Missing keys: {missing_keys}")
        if unexpected_keys:
            print(f"Control Unexpected keys: {unexpected_keys}")
        
        missing_keys, unexpected_keys = None, None
        missing_keys, unexpected_keys = self.control_depth_head.load_state_dict(self.depth_head.state_dict(), strict=False)
        if missing_keys:
            print(f"Control Depth Head Missing keys: {missing_keys}")
        if unexpected_keys:
            print(f"Control Depth Head Unexpected keys: {unexpected_keys}")